import cv2
import numpy as np
import os
import os.path as osp

IMG_SIZE = 3000


def run(DATADIR, img_dir, txt_dir, out_dir, color):
    IMAGEDIR = os.path.join(DATADIR, img_dir)
    LABELDIR = os.path.join(DATADIR, txt_dir)

    OUTDIR = os.path.join(DATADIR, out_dir)
    if not os.path.exists(OUTDIR):
        os.makedirs(OUTDIR)


    tif_names = set([name[:-4] for name in os.listdir(IMAGEDIR) if name.endswith('.tif')])
    txt_names = set([name[:-4] for name in os.listdir(LABELDIR) if name.endswith('.txt')])
    intersection_names = tif_names & txt_names

    for name in tif_names:
        name = name.split(".")[0]
        image = cv2.imread(osp.join(IMAGEDIR, name+".tif"))
        
        if name not in intersection_names:
            cv2.imwrite(osp.join(OUTDIR, name+".tif"), image)
            continue
            
        with open(osp.join(LABELDIR, name+".txt"), 'r') as fp:
            lines = fp.readlines()

        for line in lines:
            line = line.split(" ")
            label = float(line[0])
            cx = IMG_SIZE * float(line[1])
            cy = IMG_SIZE * float(line[2])
            w = IMG_SIZE * float(line[3])
            h = IMG_SIZE * float(line[4])
            angle = int(line[5])
            arg = ((cx, cy), (w, h), (angle))
            box = cv2.boxPoints(arg)
            box = np.int0(box)
            cv2.drawContours(image, [box], 0, color, 3)
            # cv2.imshow(" ", image)

        cv2.imwrite(osp.join(OUTDIR, name+".tif"), image)
        print(name + ", done...")


if __name__ == '__main__':
    DATADIR = os.path.abspath(".")
    img_dir = 'output-conf0.1-best-BiCnnNew'
    txt_dir = 'pred_txt_labels_conf0.5-1.0'
    out_dir = 'watch_output_conf0.5-1.0'
    run(DATADIR, img_dir, txt_dir, out_dir, color=(250, 200, 250))
